#!/usr/bin/env python
# _*_ coding:utf-8 _*_
#
# Copyright(c) http://m.ecpark.cn/
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
try:
    try:
        from pymongo import MongoClient as Connection
        from pymongo.errors import DuplicateKeyError, OperationFailure
        from pymongo import ASCENDING, DESCENDING, HASHED
    except ImportError:
        from pymongo import Connection
    from pymongo.errors import ConnectionFailure, AutoReconnect, DuplicateKeyError
except ImportError:
    raise ImportError("Can't import pymongo. "
                      "See http://api.mongodb.org/python/current/ "
                      "for instructions on how to install pymongo.")

import logging

log = logging.getLogger(__name__)

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
                    datefmt='%a, %d %b %Y %H:%M:%S')

BTREE = 1
HASH = 2


def enable_sharding(mongo_connect_instance, dbname):
    """
    对某个库进行enable sharding.
    :param mongo_connect_instance: pymongo.mongo_client.MongoClient 连接实例
    :param dbname: 库名
    :return:
    """
    try:
        mongo_connect_instance.admin.command('enableSharding', dbname)
    except Exception as e:
        raise ("enable sharding error on {}".format(dbname))
    return


def shard_collection(mongo_connect_instance, namespace, key):
    try:
        mongo_connect_instance.admin.command('shardcollection', namespace, key=key)
    except Exception as e:
        raise ("shard collection error on {} with{}".format(namespace, str(e)))


def get_namespace(mongo_connect_instance, namespace):
    """
    返回数据库对象属性: database 或者 collection.
    :param mongo_connect_instance:
    :param namespace:
    :return:
    """
    database_collection = namespace.split(".")
    attribute = getattr(mongo_connect_instance, database_collection[0])
    for i in range(1, len(database_collection)):
        attribute = getattr(attribute, database_collection[i])
    return attribute


def create_index(mongo_collection=None, index_keys=None, index_type="btree", is_unique=False, expire_seconds=None):
    # mongo_collection.create_index([("mike", HASHED)], background=True, unique=is_unique)
    if index_type == BTREE:
        py_index_type = ASCENDING
    elif index_type == HASH:
        if len(index_keys) != 1:
            log.info("Hash key specific more than 1 field.".format(str(index_keys)))
        py_index_type = HASHED
    else:
        log.debug(index_type)
        raise TypeError("index_type error: {}".format(str(index_type)))
    index_list = [(iter_var, py_index_type) for iter_var in index_keys]
    log.info("Now create index: {}".format(str(index_list)))
    if expire_seconds is not None:
        try:
            mongo_collection.create_index(index_list, background=True, unique=is_unique,
                                          expireAfterSeconds=expire_seconds)
        except OperationFailure as e:
            log.info(str(e))
    else:
        try:
            mongo_collection.create_index(index_list, background=True, unique=is_unique)
        except OperationFailure as e:
            log.info(str(e))


def get_current_op(mongo_connect_instance):
    return mongo_connect_instance.admin.current_op(include_all=True)


def get_all_database(mongo_connect_instance):
    return mongo_connect_instance.database_names()


def get_all_collection(mongo_database, all=1):
    if all == 1:
        db_collection = mongo_database.list_collection_names()
        # coll_filter = {"name": {"$regex": r"^(?!system\.)"}}
        # db_collection = mongo_database.list_collection_names(filter=coll_filter)
    else:
        db_collection = mongo_database.list_collection_names()
    return db_collection


if __name__ == "__main__":
    create_index(mongo_collection=None, index_keys=["a", "b"], index_type="btree", is_unique=False)
